# 垃圾分类模型核心代码
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse
import numpy as np
import time
from PIL import Image
from concurrent.futures import ThreadPoolExecutor

from cv2 import VideoCapture
import cv2
# import tensorflow as tf # TF2
import tflite_runtime.interpreter as tflite
import threading
import time
import os
import pygame
from playsound import playsound
data=[0,0,0,0]
lock = threading.Lock()
THRESH_HOLD=0.4

pool = ThreadPoolExecutor(max_workers=1)
pool_volume=ThreadPoolExecutor(max_workers=2)
def audio_play(src):
    print(src)
    pygame.mixer.init()
    pygame.mixer.music.load(src)
    pygame.mixer.music.play(0)
    time.sleep(3)
    pygame.mixer.music.unload()

    pygame.mixer.quit()

def load_labels(filename):
    with open(filename, 'r') as f:
        return [line.strip() for line in f.readlines()]
def reset(args):
        time.sleep(16)
        if args==0:


            os.system("./gpio 498 0")
        elif args==1:
            os.system("./gpio 499 0")
        elif args==2:
            os.system("./gpio 432 0")
        elif args==3:
            os.system("./gpio 434 0")

        
def upload():
    global data
    while True:
        if max(data)>20:
            index=data.index(max(data))
            if index==0:
                print("可回收物打开")
                pygame.mixer.init()

                # sound1=pygame.mixer.Sound("assets/TES.mp3")
                pool_volume.submit(audio_play,"assets/recyclable.wav")
                
                # sound1.play()

                pool.submit(reset, 0)

                os.system("./gpio 498 1")
                # pool_volume.submit(audio_play,"assets/recyclable.wav")
                
            elif index==1:
                print("其他垃圾打开")
                os.system("./gpio 499 1")
                pool.submit(reset, 1)

                pool_volume.submit(audio_play,"assets/residual.wav")

                # playsound("assets/residual.wav",True)

            elif index==2:
                print("有害垃圾打开")
                os.system("./gpio 432 1")

                pool.submit(reset, 2)

                pool_volume.submit(audio_play,"assets/hazardous.wav")


                # playsound("assets/hazardous.wav",True)

            elif index==3:
                print("厨余垃圾打开")
                os.system("./gpio 434 1")
                pool.submit(reset,3)

                pool_volume.submit(audio_play,"assets/food.wav")
                # playsound("assets/food.wav",True)
            lock.acquire()

            data=[0,0,0,0]
            lock.release()

            time.sleep(5)


           

if __name__ == '__main__':
    os.system("./gpio 498 0")
    os.system("./gpio 499 0")
    os.system("./gpio 432 0")
    os.system("./gpio 434 0")
    interpreter = tflite.Interpreter(model_path="./finalV9/model.tflite")
    labels = load_labels("./finalV9/labels.txt")

    interpreter.allocate_tensors()
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    # check the type of the input tensor
    floating_model = input_details[0]['dtype'] == np.float32
    # os.system(" aplay -Dplughw:2,0 assets/ondevice.wav")
    # playsound("assets/ondevice.wav",True)
    # sound = pygame.mixer.Sound("assets/ondevice.wav")
    # sound.play()
    # time.sleep(6)
    pygame.mixer.init()
    pygame.mixer.music.load("assets/ondevice.wav")
    pygame.mixer.music.play(0)
    time.sleep(5)
    pygame.mixer.music.unload()

    pygame.mixer.quit()

    # NxHxWxC, H:1, W:2
    height = input_details[0]['shape'][1]
    width = input_details[0]['shape'][2]
    cap = VideoCapture(-1)

    cap.set(10, 100)  # 亮度
    t1=threading.Thread(target=upload,)
    t1.start()  
    while True:
        a=(time.process_time())

        success, img = cap.read()
 #       cv2.imshow("dd",img)
        img_origin=img
#        cv2.waitKey(1)
        img= Image.fromarray(img).convert('RGB').resize((width,height))
        # add N dim
        input_data = np.expand_dims(img, axis=0)

        if floating_model:
            input_data = (np.float32(input_data) - 127.5) / 127.5
        interpreter.set_tensor(input_details[0]['index'], input_data)

        interpreter.invoke()

        output_data = interpreter.get_tensor(output_details[0]['index'])
        results = np.squeeze(output_data)
	
        top_k = results.argsort()[-5:][::-1]
        cn=labels[top_k[0]].split()[0]
        can_class=cn.split(".")[0]
        en=labels[top_k[0]].split()[1]
        probability=results[top_k[0]]/255.0
        if probability>=THRESH_HOLD:
            if "可回收物" in cn:
                data[0]+=1
            elif "其他垃圾" in cn:
                data[1]+=1
            elif "有害垃圾" in cn:
                data[2]+=1
            elif "厨余垃圾" in cn:
                data[3]+=1
                               
        if floating_model:
            # print('{:08.6f}: {}'.format(float(results[top_k[0]]), labels[top_k[0]]))
            print('{:08.6f}: {}'.format(float(results[top_k[0]]), cn))
        else:
            # print('{:08.6f}: {}'.format(float(results[top_k[0]] / 255.0), labels[top_k[0]]))
            print('{:08.6f}: {}'.format(float(results[top_k[0]] / 255.0), cn))
        print(time.process_time()-a)
        cv2.putText(img_origin, en, org=(10, 25), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=1,color=(255, 0, 0), thickness=2)

        cv2.imshow("dd",img_origin)
        cv2.waitKey(1)


        
